from pgmpy.models import BayesianNetwork
from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import VariableElimination

class ReasoningSpace:
    """Reasoning Space implementation (Section 3.3)"""
    def __init__(self):
        self.memory = []
        self.reasoning_weights = {
            "color_harmony": 0.25,
            "composition": 0.25,
            "style_coherence": 0.3,
            "trend_alignment": 0.2
        }
        
        # Business impact factors
        self.business_factors = {
            "purchase_likelihood": 0.4,
            "customer_engagement": 0.3,
            "trend_relevance": 0.3
        }

        # Initialize Bayesian Network
        self.network = BayesianNetwork([
            ('task_success', 'action_choice'),
            ('resource_available', 'task_success'),
            ('environmental_condition', 'task_success'),
            ('q_value', 'action_choice'),
            ('uncertainty', 'task_success')
        ])
        
        self._initialize_cpds()
    
    def _initialize_cpds(self):
        """Initialize Conditional Probability Distributions"""
        cpd_resource = TabularCPD(
            variable='resource_available',
            variable_card=2,
            values=[[0.7], [0.3]]
        )
        
        cpd_environment = TabularCPD(
            variable='environmental_condition',
            variable_card=2,
            values=[[0.8], [0.2]]
        )
        
        cpd_uncertainty = TabularCPD(
            variable='uncertainty',
            variable_card=2,
            values=[[0.6], [0.4]]
        )
        
        cpd_success = TabularCPD(
            variable='task_success',
            variable_card=2,
            evidence=['resource_available', 'environmental_condition', 'uncertainty'],
            evidence_card=[2, 2, 2],
            values=[
                [0.9, 0.7, 0.7, 0.5, 0.6, 0.3, 0.2, 0.1],
                [0.1, 0.3, 0.3, 0.5, 0.4, 0.7, 0.8, 0.9]
            ]
        )
        
        # Add CPD for action_choice
        cpd_action = TabularCPD(
            variable='action_choice',
            variable_card=2,  # Binary: Take action/Wait
            evidence=['task_success', 'q_value'],
            evidence_card=[2, 2],
            values=[
                [0.9, 0.7, 0.3, 0.1],  # Take action probabilities
                [0.1, 0.3, 0.7, 0.9]   # Wait probabilities
            ]
        )
        
        # CPD for q_value
        cpd_qvalue = TabularCPD(
            variable='q_value',
            variable_card=2,  # Binary: High/Low
            values=[[0.5], [0.5]]  # Initially uniform
        )
        
        self.network.add_cpds(
            cpd_resource, cpd_environment, 
            cpd_uncertainty, cpd_success,
            cpd_action, cpd_qvalue
        )
    
    def decide(self, perception, plan):
        """Make decision based on perception and plan"""
        inference = VariableElimination(self.network)
        
        # Set up evidence based on perception and plan
        evidence = {
            'resource_available': 1 if perception['energy'] > 50 else 0,
            'environmental_condition': 1 if perception['knowledge'] > 50 else 0,
            'uncertainty': 0 if len(plan['primitive_tasks']) < 5 else 1
        }
        
        # Query the network
        result = inference.query(variables=['task_success'], evidence=evidence)
        
        return {
            'action': 'proceed' if result.values[1] > 0.5 else 'wait',
            'confidence': result.values[1]
        } 

    def evaluate_beauty(self, perception, plan):
        """Evaluate beauty based on multiple aspects and their relationships
        
        Args:
            perception: Visual features from perception space
            plan: Analysis plan from planning space
            
        Returns:
            dict: Beauty evaluation including:
                - overall_score: Combined beauty score
                - aspect_scores: Individual aspect scores
                - business_impact: Estimated business impact
                - confidence: Confidence in evaluation
        """
        # Initialize results
        aspect_scores = {}
        confidence_scores = {}
        
        # Evaluate each aspect according to plan
        for step in plan["steps"]:
            aspect = step["type"]
            if aspect == "color_harmony":
                score, conf = self._evaluate_color_harmony(perception)
            elif aspect == "composition":
                score, conf = self._evaluate_composition(perception)
            elif aspect == "style_coherence":
                score, conf = self._evaluate_style_coherence(perception)
            elif aspect == "trend_alignment":
                score, conf = self._evaluate_trend_alignment(perception)
            else:
                continue
                
            aspect_scores[aspect] = score
            confidence_scores[aspect] = conf
        
        # Calculate overall beauty score
        overall_score = sum(
            score * self.reasoning_weights[aspect]
            for aspect, score in aspect_scores.items()
        )
        
        # Calculate business impact
        business_impact = self._calculate_business_impact(aspect_scores)
        
        # Calculate overall confidence
        overall_confidence = sum(
            conf * self.reasoning_weights[aspect]
            for aspect, conf in confidence_scores.items()
        )
        
        result = {
            "overall_score": overall_score,
            "aspect_scores": aspect_scores,
            "business_impact": business_impact,
            "confidence": overall_confidence
        }
        
        # Store in memory
        self.memory.append(result)
        
        return result

    def _evaluate_color_harmony(self, perception):
        """Evaluate color harmony and its impact"""
        if "color_features" not in perception:
            return 0.0, 0.0
            
        color_features = perception["color_features"]
        harmony_score = color_features.get("harmony_score", 0.0)
        confidence = 0.8  # High confidence in color analysis
        
        return harmony_score, confidence

    def _evaluate_composition(self, perception):
        """Evaluate composition quality"""
        if "spatial_features" not in perception:
            return 0.0, 0.0
            
        spatial_features = perception["spatial_features"]
        composition_score = spatial_features.get("composition", 0.0)
        confidence = 0.7  # Medium-high confidence
        
        return composition_score, confidence

    def _evaluate_style_coherence(self, perception):
        """Evaluate style coherence and appeal"""
        if "semantic_features" not in perception:
            return 0.0, 0.0
            
        semantic_features = perception["semantic_features"]
        style_score = 0.0
        confidence = 0.6  # Medium confidence due to subjectivity
        
        # TODO: Implement style coherence evaluation
        return style_score, confidence

    def _evaluate_trend_alignment(self, perception):
        """Evaluate alignment with current trends"""
        if "semantic_features" not in perception:
            return 0.0, 0.0
            
        semantic_features = perception["semantic_features"]
        trend_score = 0.0
        confidence = 0.5  # Lower confidence due to trend volatility
        
        # TODO: Implement trend alignment evaluation
        return trend_score, confidence

    def _calculate_business_impact(self, aspect_scores):
        """Calculate potential business impact"""
        return {
            "purchase_likelihood": self._estimate_purchase_likelihood(aspect_scores),
            "customer_engagement": self._estimate_engagement(aspect_scores),
            "trend_relevance": aspect_scores.get("trend_alignment", 0.0)
        }

    def _estimate_purchase_likelihood(self, scores):
        """Estimate likelihood of purchase based on scores"""
        return (
            scores.get("style_coherence", 0.0) * 0.4 +
            scores.get("trend_alignment", 0.0) * 0.3 +
            scores.get("color_harmony", 0.0) * 0.3
        )

    def _estimate_engagement(self, scores):
        """Estimate potential customer engagement"""
        return (
            scores.get("composition", 0.0) * 0.4 +
            scores.get("color_harmony", 0.0) * 0.3 +
            scores.get("style_coherence", 0.0) * 0.3
        ) 